"""
佑铭方案
疫情+通用分别8:2切分，合并
80%在切5份
"""
import os
import json
import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold, train_test_split


SEED = 42

# load data
with open(r'../raw/train/usual_train.txt', 'r', encoding='utf-8') as f:
    usual_data = json.load(f)
usual_content = [i['content'] for i in usual_data]
usual_content = np.array(usual_content)
usual_labels = [i['label'] for i in usual_data]
usual_labels = np.array(usual_labels)
with open(r'../raw/train/virus_train.txt', 'r', encoding='utf-8') as f:
    virus_data = json.load(f)
virus_content = [i['content'] for i in virus_data]
virus_content = np.array(virus_content)
virus_labels = [i['label'] for i in virus_data]
virus_labels = np.array(virus_labels)

# 20%
usual_content, usual_dev_content, usual_labels, usual_dev_labels = train_test_split(usual_content,
                                                                                    usual_labels,
                                                                                    stratify=usual_labels,
                                                                                    test_size=0.2,
                                                                                    shuffle=True,
                                                                                    random_state=SEED)
virus_content, virus_dev_content, virus_labels, virus_dev_labels = train_test_split(virus_content,
                                                                                    virus_labels,
                                                                                    stratify=virus_labels,
                                                                                    test_size=0.2,
                                                                                    shuffle=True,
                                                                                    random_state=SEED)

# 20%当作无标签
content_dev = np.concatenate([usual_dev_content, virus_dev_content])
labels_dev = np.concatenate([usual_dev_labels, virus_dev_labels])
df = pd.DataFrame()
df['content'] = content_dev
df['labels'] = labels_dev
df.to_csv(r'./test.csv', index=0, encoding='utf-8')


# 80%分五折
content = np.concatenate([usual_content, virus_content])
labels = np.concatenate([usual_labels, virus_labels])
kfold = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED)
# usual data
for index, (train_index, dev_index) in enumerate(kfold.split(content, labels)):
    print('{}/{}'.format(index+1, 5))
    # train
    train_X = content[train_index]
    train_Y = labels[train_index]
    df = pd.DataFrame()
    df['content'] = train_X
    df['labels'] = train_Y
    df.to_csv(r'./train{}.csv'.format(index), index=0, encoding='utf-8')

    # dev
    dev_X = content[dev_index]
    dev_Y = labels[dev_index]
    df = pd.DataFrame()
    df['content'] = dev_X
    df['labels'] = dev_Y
    df.to_csv(r'./dev{}.csv'.format(index), index=0, encoding='utf-8')
